package com.easy.slurm.slurm.api.Impl;

import com.easy.slurm.exception.BussinessException;
import com.easy.slurm.exception.ResultCode;
import com.easy.slurm.slurm.Jsch.JschClient;
import com.easy.slurm.slurm.api.SlurmApi;
import com.easy.slurm.slurm.bean.*;
import com.easy.slurm.slurm.config.SlurmConstants;
import com.easy.slurm.vo.ResultVo;
import com.jcraft.jsch.JSchException;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @author bing.bai
 * @create 2025/6/4
 */
@Service
@Slf4j
public class SlurmApiImpl implements SlurmApi {

    private static String SINFO = "sinfo -o \"%P|%a|%l|%D|%t|%N\" -h";

    private static String SBATCH = "sbatch";

    private static String SQUEUE = "squeue --format=\"%i|%P|%j|%u|%T|%M|%l|%D|%C|%R\" -h";

    private static String SCANCEL = "scancel ";

    private static String SCONTROL = "scontrol show job ";


    @Override
    public List<SInfo> sinfo() throws JSchException, IOException {
        ArrayList<SInfo> sinfoList = new ArrayList();
        JschClient instance = JschClient.getInstance();
        JschClient.RunCommandResult runCommandResult = instance.runCommand(SINFO);
        if (runCommandResult.getExitStatus() == 0) {
            List<String> stdout = runCommandResult.getStdout();
            for (int i = 0; i < stdout.size(); i++) {
                SInfo sInfo = new SInfo();
                String line = stdout.get(i);
                String[] snode = line.split("\\|");
                sInfo.setPartition(snode[0]);
                sInfo.setAvail(snode[1]);
                sInfo.setTimeLimit(snode[2]);
                sInfo.setNodes(Integer.parseInt(snode[3]));
                sInfo.setState(snode[4]);
                sInfo.setNodeList(snode[5]);
                sinfoList.add(sInfo);
            }
        } else {
            throw new BussinessException(runCommandResult.getErrorMsg());
        }
        return sinfoList;
    }

    @Override
    public ResultVo submitSbatchJob(SbatchJobSubmit sbatchJobSubmit) throws JSchException, IOException {
        if (StringUtils.isBlank(sbatchJobSubmit.getFileName())) {
            throw new BussinessException(ResultCode.FILENAMENOTEXIST);
        }
        String sbatchCommand = combineSbatchJobCommand(sbatchJobSubmit);
        JschClient instance = JschClient.getInstance();
        JschClient.RunCommandResult runCommandResult = instance.runCommand(sbatchCommand);
        if (runCommandResult.isSuccess()) {
            return ResultVo.success(runCommandResult.getSuccessMsg());
        } else {
            return ResultVo.fail(runCommandResult.getErrorMsg());
        }
    }

    @Override
    public List<Squeue> squeue() throws JSchException, IOException {
        ArrayList<Squeue> squeues = new ArrayList<>();
        JschClient instance = JschClient.getInstance();
        JschClient.RunCommandResult runCommandResult = instance.runCommand(SQUEUE);
        if (runCommandResult.isSuccess()) {
            List<String> stdout = runCommandResult.getStdout();
            for (int i = 0; i < stdout.size(); i++) {
                Squeue squeue = new Squeue();
                String line = stdout.get(i);
                String[] job = line.split("\\|");
                squeue.setJobId(Integer.parseInt(job[0]));
                squeue.setPartition(job[1]);
                squeue.setName(job[2]);
                squeue.setUser(job[3]);
                squeue.setStatus(job[4]);
                squeue.setTime(job[5]);
                squeue.setTimeLimited(job[6]);
                squeue.setNodes(Integer.parseInt(job[7]));
                squeue.setCpus(Integer.parseInt(job[8]));
                squeue.setNodeList_reason(job[9]);
                squeues.add(squeue);
            }
        }
        return squeues;
    }

    @Override
    public ResultVo cancel(CancelJobRequest cancelJobRequest) throws JSchException, IOException {
        if (cancelJobRequest.getJobIdList().isEmpty()) {
            throw new BussinessException(ResultCode.JOBIDNOTEXIST);
        }
        List<Integer> jobIdList = cancelJobRequest.getJobIdList();
        String command = SCANCEL + jobIdList.stream()
                .map(Object::toString)
                .collect(Collectors.joining(" "));
        JschClient instance = JschClient.getInstance();
        JschClient.RunCommandResult runCommandResult = instance.runCommand(command);
        if (runCommandResult.isSuccess()) {
            return ResultVo.success(runCommandResult.getSuccessMsg());
        } else {
            return ResultVo.fail(runCommandResult.getErrorMsg());
        }
    }

    @Override
    public String scontrol(String jobId) throws JSchException, IOException {
        JschClient instance = JschClient.getInstance();
        JschClient.RunCommandResult runCommandResult = instance.runCommand(SCONTROL + jobId);
        if (runCommandResult.isSuccess()) {
            return runCommandResult.getSuccessMsg();
        } else {
            return runCommandResult.getErrorMsg();
        }
    }

    @Override
    public ResultVo uploadFile(FileUpload fileUpload) {
        JschClient instance = JschClient.getInstance();
        instance.uploadFile(fileUpload.getFile(), SlurmConstants.WORKDIR + "/" + fileUpload.getTenantId(), fileUpload.getFileName());
        return ResultVo.success(ResultCode.SUCCESS);
    }


    private String combineSbatchJobCommand(SbatchJobSubmit sbatchJobSubmit) {
        StringBuilder command = new StringBuilder(SBATCH);
        command.append(" -D " + SlurmConstants.WORKDIR + "/" + sbatchJobSubmit.getTenantId());
        if (StringUtils.isNotBlank(sbatchJobSubmit.getJobName())) {
            command.append(" -J " + sbatchJobSubmit.getJobName());
        }
        if (StringUtils.isNotBlank(sbatchJobSubmit.getPartition())) {
            command.append(" -p " + sbatchJobSubmit.getPartition());
        }
        if (StringUtils.isNotBlank(sbatchJobSubmit.getTime())) {
            command.append(" -t " + sbatchJobSubmit.getTime());
        }
        if (sbatchJobSubmit.getNodes() != null) {
            command.append(" --nodes=" + sbatchJobSubmit.getNodes());
        }
        if (sbatchJobSubmit.getNtasks() != null) {
            command.append(" --ntasks=" + sbatchJobSubmit.getNtasks());
        }
        if (sbatchJobSubmit.getCpusPerTas() != null) {
            command.append(" --cpus-per-task=" + sbatchJobSubmit.getCpusPerTas());
        }
        if (StringUtils.isNotBlank(sbatchJobSubmit.getMem())) {
            command.append(" --mem= " + sbatchJobSubmit.getMem());
        }
        if (StringUtils.isNotBlank(sbatchJobSubmit.getMemPerCpu())) {
            command.append(" --mem-per-cpu= " + sbatchJobSubmit.getMemPerCpu());
        }
        if (StringUtils.isNotBlank(sbatchJobSubmit.getGres())) {
            command.append(" --gres= " + sbatchJobSubmit.getGres());
        }
        if (sbatchJobSubmit.getExclusive() != null && sbatchJobSubmit.getExclusive()) {
            command.append(" --exclusive= ");
        }
        if (StringUtils.isNotBlank(sbatchJobSubmit.getNodelist())) {
            command.append(" --nodelist= ");
        }
        if (StringUtils.isNotBlank(sbatchJobSubmit.getExclude())) {
            command.append(" --exclude= " + sbatchJobSubmit.getExclude());
        }
        command.append(" " + sbatchJobSubmit.getFileName());
        return command.toString();
    }


}
